import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from matplotlib.animation import FuncAnimation
from heapq import heappush, heappop

class MapPlot:
    def __init__(self, grid_size=(10, 10), obstacles=None):
        """
        初始化自动泊车系统
        
        Parameters:
        -----------
        grid_size : tuple, 网格大小 (width, height)
        obstacles : list, 障碍物坐标列表 [(x1, y1), (x2, y2), ...]
        """
        self.grid_size = grid_size
        self.obstacles = obstacles if obstacles else []
        
        # 图形对象
        self.fig, self.ax = None, None
        self.car = None
        self.path_line = None
        self.ani = None
        
        # 状态变量
        self.start = None
        self.goal = None
        self.path = None
        self.current_position = None
        
        # 初始化图形
        self._init_plot()
    
    def _init_plot(self):
        """初始化绘图环境"""
        self.fig, self.ax = plt.subplots(figsize=(8, 8))
        self.ax.set_xlim(0, self.grid_size[0])
        self.ax.set_ylim(0, self.grid_size[1])
        self.ax.set_aspect('equal')
        self.ax.grid(True)
        self.ax.set_title("Automated Parking Valet Simulation")
        
        # 绘制障碍物
        for obs in self.obstacles:
            self.ax.add_patch(patches.Rectangle(
                (obs[0]-0.5, obs[1]-0.5), 1, 1, color='red', alpha=0.7))
        plt.ion()
    def set_ego_car_position(self, start_pos, goal_pos, path=None):
        """
        设置车辆起始位置和目标位置，并绘制路径（可选）
        
        Parameters:
        -----------
        start_pos : tuple, 起始位置 (x, y)
        goal_pos : tuple, 目标位置 (x, y)
        path : list, 可选，路径坐标列表 [(x1,y1), (x2,y2), ...]
        
        Returns:
        --------
        bool: 是否成功设置
        """
        self.start = start_pos
        self.goal = goal_pos
        self.current_position = start_pos
        self.path = path
        
        # 清除之前的图形元素
        if self.path_line:
            self.path_line.remove()
            self.path_line = None
        if self.car:
            self.car.remove()
            self.car = None
        
        # 清除之前的起点终点标记
        if hasattr(self, 'start_marker'):
            self.start_marker.remove()
        if hasattr(self, 'goal_marker'):
            self.goal_marker.remove()
        
        # 绘制起点和终点
        self.start_marker, = self.ax.plot(start_pos[0], start_pos[1], 'go', 
                    markersize=10, label='Start')
        self.goal_marker, = self.ax.plot(goal_pos[0], goal_pos[1], 'bo', 
                    markersize=10, label='Goal')
        
        # 如果有路径，绘制路径
        if path:
            path_x = [p[0] for p in path]
            path_y = [p[1] for p in path]
            self.path_line, = self.ax.plot(path_x, path_y, 'y--', linewidth=2, label='Path')
        
        # 创建车辆
        self.car = patches.Rectangle(
            (start_pos[0]-0.3, start_pos[1]-0.2), 0.6, 0.4, 
            color='blue', alpha=0.8, label='Ego Car')
        self.ax.add_patch(self.car)
        
        # 添加或更新图例
        if self.ax.legend_:
            self.ax.legend_.remove()
        self.ax.legend(loc='upper right')
        
        # 刷新图形
        self.fig.canvas.draw_idle()
        
        return True

    def update_path(self, new_path):
        """
        更新路径显示
        
        Parameters:
        -----------
        new_path : list, 新的路径坐标列表
        """
        self.path = new_path
        
        # 清除旧路径
        if self.path_line:
            self.path_line.remove()
            self.path_line = None
        
        # 绘制新路径
        if new_path is not None and len(new_path) > 0:
            # 转换为NumPy数组以确保一致性
            path_array = np.array(new_path)
            
            # 提取x, y坐标
            path_x = path_array[:, 0]
            path_y = path_array[:, 1]
            
            self.path_line, = self.ax.plot(path_x, path_y, 'y--', linewidth=2, label='Path')
        
        # 更新图例和刷新图形
        if self.ax.legend_:
            self.ax.legend_.remove()
        self.ax.legend(loc='upper right')
        self.fig.canvas.draw_idle()

    def clear_path(self):
        """
        清除路径显示
        """
        if self.path_line:
            self.path_line.remove()
            self.path_line = None
            self.path = None
            self.fig.canvas.draw_idle()
    
    
    
    def update_car_position(self, new_position):
        """
        更新车辆位置
        
        Parameters:
        -----------
        new_position : tuple, 新的车辆位置 (x, y)
        """
        if self.car:
            self.car.set_xy((new_position[0]-0.3, new_position[1]-0.2))
            self.current_position = new_position
            self.fig.canvas.draw_idle()  # 立即更新图形
    
    def animate_path(self, interval=500):
        """
        动画展示路径跟踪
        
        Parameters:
        -----------
        interval : int, 动画帧间隔(毫秒)
        """
        if not self.path:
            print("错误：请先设置车辆位置和路径！")
            return
        
        def update(frame):
            if frame < len(self.path):
                x, y = self.path[frame]
                self.car.set_xy((x-0.3, y-0.2))
            return self.car,
        
        self.ani = FuncAnimation(
            self.fig, update, frames=len(self.path)+10, 
            interval=interval, blit=True, repeat=False
        )
        
        plt.show()
    
    def get_path(self):
        """获取当前路径"""
        return self.path
    
    def get_current_position(self):
        """获取当前车辆位置"""
        return self.current_position
    
    def add_obstacle(self, obstacle_pos):
        """
        添加障碍物
        
        Parameters:
        -----------
        obstacle_pos : tuple, 障碍物坐标 (x, y)
        """
        if obstacle_pos not in self.obstacles:
            self.obstacles.append(obstacle_pos)
            # 重新绘制障碍物
            self.ax.add_patch(patches.Rectangle(
                (obstacle_pos[0]-0.5, obstacle_pos[1]-0.5), 1, 1, color='red', alpha=0.7))
            self.fig.canvas.draw_idle()
    
    def remove_obstacle(self, obstacle_pos):
        """
        移除障碍物
        
        Parameters:
        -----------
        obstacle_pos : tuple, 障碍物坐标 (x, y)
        """
        if obstacle_pos in self.obstacles:
            self.obstacles.remove(obstacle_pos)
            # 需要重新绘制整个图形来移除障碍物
            self._redraw_plot()
    
    def _redraw_plot(self):
        """重新绘制整个图形"""
        self.ax.clear()
        self._init_plot()
        # 重新设置车辆位置（如果已设置）
        if self.start and self.goal:
            self.set_ego_car_position(self.start, self.goal)
    
    def show(self):
        """显示图形"""
        plt.show()
    
    def close(self):
        """关闭图形"""
        plt.close(self.fig)


# 使用示例
if __name__ == "__main__":
    # 定义障碍物
    obstacles = [(2, 2), (2, 3), (3, 3), (4, 4), (5, 5), (7, 7), (7, 8)]
    
    # 创建自动泊车系统
    parking_system = MapPlot(grid_size=(10, 10), obstacles=obstacles)
    
    success = parking_system.set_ego_car_position(start_pos=(1, 1), goal_pos=(8, 8))
    
    # 方法1：直接显示静态图
    # parking_system.show()
    
    # 方法2：启动动画（取消注释以下行）
    # parking_system.animate_path(interval=300)
    
    # 方法3：手动更新车辆位置示例
    for i in range(5):
        new_pos = (i+2, i+2)
        parking_system.update_car_position(new_pos)
        plt.pause(1)  # 暂停1秒看效果